# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
MCTS 无代码模式实现

本模块实现了基于蒙特卡洛树搜索的无代码推理算法，用于解决数学推理问题。
与标准 MCTS 不同，此版本移除了代码生成和执行部分，专注于纯推理思考过程。

主要功能:
- 实现无代码 MCTS 算法的核心流程
- 支持节点选择和扩展策略（继承自标准 MCTS）
- 提供价值估计和反向传播机制
- 移除代码执行和错误处理
- 支持中间结果记录和评估

核心组件:
- MCTSWoCode: 无代码 MCTS 算法实现类
- 简化的节点创建: 移除代码相关状态处理
- 优化的解析逻辑: 专注于思考步骤和最终答案
- 兼容的接口设计: 保持与原有系统的接口一致性

关键差异:
- 使用 rstar_step_result_unwrap_wo_code 替代原有解析函数
- 移除代码执行相关的错误处理逻辑
- 保持所有其他功能完整性，包括多模态支持
"""
from __future__ import annotations  # 启用延迟注解，支持前向引用类型提示
from typing import Dict, Any, Optional, Type, List, Tuple  # 类型提示工具，提供静态类型检查支持
from pydantic import field_validator  # Pydantic 字段验证装饰器，用于配置参数验证
from rstar_deepthink.agents.mcts import MCTS  # 继承的标准 MCTS 类
from rstar_deepthink.agents.utils import rstar_step_result_unwrap_wo_code  # 无代码模式解析函数
from rstar_deepthink.nodes import MCTSNode  # MCTS 节点类
from rstar_deepthink.constants import (  # 项目常量定义
    TOO_MANY_STEPS,       # 步数超限的终止标记
    NO_VALID_CHILD,       # 无有效子节点的终止标记
)


class MCTSWoCode(MCTS):
    """
    无代码模式的蒙特卡洛树搜索（MCTS）算法实现类
    
    继承自标准 MCTS 类，实现了无代码版本的 MCTS 算法流程。
    主要差异在于移除了代码生成、执行和相关错误处理，专注于纯思考推理。
    
    核心特点:
    1. 继承标准 MCTS 的所有核心功能（选择、评估、回传等）
    2. 使用专门的无代码解析函数处理 LLM 输出
    3. 移除代码执行相关的状态和错误处理
    4. 保持与原有系统的完全兼容性
    5. 支持所有现有的多模态和 action 机制
    
    算法流程:
    1. 选择（Selection）：完全复用父类实现
    2. 扩展（Expansion）：使用简化的节点创建逻辑
    3. 模拟（Simulation）：完全复用父类实现
    4. 反向传播（Backpropagation）：完全复用父类实现
    """
    
    def __init__(self, **kwargs):
        """
        初始化无代码 MCTS 实例
        
        设置无代码模式的解析函数，其他配置完全继承父类。
        这确保了除解析逻辑外的所有功能都保持一致。
        
        参数:
            **kwargs: 传递给父类的初始化参数
        """
        # 调用父类初始化，保持所有原有功能
        super().__init__(**kwargs)
        
        # 替换为无代码模式的解析函数
        self.step_unwrap = rstar_step_result_unwrap_wo_code
    
    @field_validator("config")
    def validate_config(cls, cfg: Any):
        """
        验证无代码 MCTS 配置参数
        
        确保配置参数符合无代码 MCTS 算法的要求，主要验证模式设置。
        
        参数:
            cfg: 配置对象
            
        返回:
            验证后的配置对象
            
        异常:
            ValueError: 当配置参数不符合要求时抛出
        """
        # 先调用基类的验证（但跳过模式检查）
        from rstar_deepthink.agents.tree import BaseTree
        BaseTree.validate_config(cfg)
        
        # 验证无代码模式
        if not cfg.mode == "mcts_wo_code":
            raise ValueError(f"Wrong value for config mode, expected 'mcts_wo_code', got '{cfg.mode}'.")
        if cfg.stop is None:
            raise ValueError(f"Wrong value for config stop, cannot be None")
        return cfg

    def create_child(
        self, 
        step_result: str, 
        parser_result: Dict[str, str], 
        node: Type[MCTSNode],
        idx: int,
    ) -> None:
        """
        为指定父节点创建子节点 - 无代码版本
        
        根据 LLM 生成的步骤结果和解析结果，创建新的子节点并设置其状态。
        与标准版本的主要区别是移除了代码执行相关的处理逻辑。
        
        参数:
            step_result: LLM 生成的步骤文本
            parser_result: 解析后的结构化结果
            node: 父节点
            idx: 子节点索引（当前未使用）
            
        处理逻辑:
        1. 如果解析失败 (parser_result is None)：标记为终端节点，设置 NO_VALID_CHILD
        2. 如果包含最终答案：标记为终端节点，评估答案正确性
        3. 如果包含其他 action：保持现有处理逻辑（移除代码执行部分）
        4. 普通思考文本：设置为中间节点
        5. 检查深度限制，超过最大深度则终止
        
        关键修改:
        - 移除所有代码执行相关逻辑（python_interpreter action 处理）
        - 移除代码错误处理（TOO_MANY_CODE_ERRORS）
        - 保留所有其他 action 类型的处理（确保多模态兼容性）
        - 保持深度检查和其他终止条件
        """
        # 创建新的子节点，设置基本属性
        new_node = self.create_node(parent=node)
        parent_child_count = len(node.children)
        new_node.tag = f"{node.tag}.{parent_child_count + 1}"
        new_node.depth = node.depth + 1

        # 处理解析失败的情况
        if parser_result is None:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = NO_VALID_CHILD
            self.eval_final_answer(new_node)
        # 处理包含最终答案的情况
        elif parser_result["final_answer"]:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = parser_result["final_answer"]
            self.eval_final_answer(new_node)
        # 处理包含 action 的情况（保持现有逻辑，但移除代码执行）
        elif parser_result["action"]:
            # 注意：无代码模式下，parser_result["action"] 应该为空
            # 但为了保持多模态兼容性，这里保留其他 action 的处理框架
            # 实际上，rstar_step_result_unwrap_wo_code 不会设置 action
            new_node.state["action"] = parser_result["action"]
            new_node.state["action_input"] = parser_result["action_input"]
            new_node.state["text"] = step_result
            # 移除代码执行逻辑，其他 action 类型可以在这里扩展
        else:
            # 处理普通思考文本
            new_node.state["text"] = step_result

        # 检查深度限制，超过最大深度则强制终止
        if not new_node.is_terminal and new_node.depth > self.config.max_depth:
            new_node.is_terminal = True
            new_node.state["final_answer"] = TOO_MANY_STEPS
            self.eval_final_answer(new_node)

        # 将新节点添加到父节点的子节点列表
        node.children.append(new_node) 